#include "ohos_video_h264_encoder.h"

#include <algorithm>
#include <cstdint>
#include <limits>
#include <memory>
#define TIMEOUT 8
#define THIS_FILE "OHOSMediaEncode"
#include <fstream>
#include <iostream>

OHOSMediaEncode::OHOSMediaEncode(const std::string &mime)
    : H26xEncoder(mime), _psInserter(mediastreamer::H26xToolFactory::get(mime).createParameterSetsInserter()) {
	_vsize.width = 0;
	_vsize.height = 0;
	data_ = std::make_unique<ohos::CodecData>();
	createImpl();
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
}

OHOSMediaEncode::~OHOSMediaEncode() {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	if (_impl) OH_VideoEncoder_Destroy(_impl);
}

ohos::CodecData *OHOSMediaEncode::GetCodecData() {
	return data_.get();
}

void OHOSMediaEncode::createImpl() {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	_impl = OH_VideoEncoder_CreateByMime(_mime.c_str());
	if (_impl == nullptr) {
		ms_error("errorCode %s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	}
}

int32_t OHOSMediaEncode::preConfig() {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	int32_t ret = setCallback();
	ret = OH_VideoEncoder_Prepare(_impl);
	if (ret != AV_ERR_OK) {
		ms_error("%s %s,%d errorCode", THIS_FILE, __FUNCTION__, __LINE__);
	}
	return ret;
}
int32_t OHOSMediaEncode::setCallback() {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	int32_t ret = 0;
	OH_AVCodecCallback cb = {&OnCodecError, &OnCodecFormatChange, &OnNeedInputBuffer, &OnNewOutputBuffer};
	ret = OH_VideoEncoder_RegisterCallback(_impl, cb, this);
	if (ret != AV_ERR_OK) {
		ms_error("%s %s,%d errorCode", THIS_FILE, __FUNCTION__, __LINE__);
	}
	ms_message("%s %s,%d,ret:%d", THIS_FILE, __FUNCTION__, __LINE__, ret);
	return ret;
}

void OHOSMediaEncode::setFps(float fps) {
	_fps = fps;
	if (isRunning() && _impl) {
		OH_AVFormat *avformat = OH_AVFormat_Create();
		OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_FRAME_RATE, _fps);
		int32_t ret = OH_VideoEncoder_Configure(_impl, avformat);
		if (ret != AV_ERR_OK) {
			ms_error("errorCode %s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
		}
		ms_message("OHOSMediaEncode: %s,%d", __FUNCTION__, __LINE__);
		OH_AVFormat_Destroy(avformat);
	}
}

void OHOSMediaEncode::setBitrate(int bitrate) {
	_bitrate = bitrate;
	if (isRunning() && _impl) {
		OH_AVFormat *avformat = OH_AVFormat_Create();
		OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_BITRATE, (_bitrate * 9) / 10);
		int32_t ret = OH_VideoEncoder_Configure(_impl, avformat);
		if (ret != AV_ERR_OK) {
			ms_error("errorCode %s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
		}
		ms_message("OHOSMediaEncode: %s,%d", __FUNCTION__, __LINE__);
		OH_AVFormat_Destroy(avformat);
	}
}

void OHOSMediaEncode::start() {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);

	if (_impl == nullptr) {
		ms_error("MediaCodecEncoder: errorCode starting failed. No MediaCodec instance.");
		return;
	}
	if (_isRunning) {
		ms_warning("MediaCodecEncoder: errorCode encoder already started");
		return;
	}
	configureImpl();
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	if (OH_VideoEncoder_Start(_impl) != AV_ERR_OK) {
		ms_error("%s %s,%d errorCode", THIS_FILE, __FUNCTION__, __LINE__);
	}
	_isRunning = true;
	ms_message("MediaCodecEncoder: encoder successfully started");
}

void OHOSMediaEncode::stop() {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	if (_impl == nullptr) return;
	if (!_isRunning) {
		ms_warning("MediaCodecEncoder: encoder already stopped");
		return;
	}
	if (_impl) {
		OH_VideoEncoder_Flush(_impl);
		OH_VideoEncoder_Stop(_impl);
		OH_VideoEncoder_Destroy(_impl);
		_impl = nullptr;
	}
	_psInserter->flush();
	_isRunning = false;
	_pendingFrames = 0;
	_firstImageQueued = false;
}

void OHOSMediaEncode::OnCodecError(OH_AVCodec *codec, int32_t errorCode, void *userData) {
	(void)codec;
	(void)userData;
	ms_message("%s %s,%d errorCode:%d", THIS_FILE, __FUNCTION__, __LINE__, errorCode);
}

void OHOSMediaEncode::OnCodecFormatChange(OH_AVCodec *codec, OH_AVFormat *format, void *userData) {
	// only trigger for surface mode
	(void)codec;
	if (userData == nullptr || format == nullptr) {
		return;
	}
	ms_message("OHOSMediaEncode: %s,%d", __FUNCTION__, __LINE__);
	OH_AVFormat *avformat = OH_AVFormat_Create();
	OH_AVFormat_Copy(avformat, format);
	OH_AVFormat_Destroy(avformat);
}

void OHOSMediaEncode::OnNeedInputBuffer(OH_AVCodec *codec, uint32_t index, OH_AVBuffer *buffer, void *userData) {
	(void)codec;
	if (userData == nullptr || buffer == nullptr) {
		return;
	}
	ms_message("OHOSMediaEncode: %s,%d", __FUNCTION__, __LINE__);
	OHOSMediaEncode *encoder = static_cast<OHOSMediaEncode *>(userData);
	OH_AVFormat *format = OH_VideoEncoder_GetInputDescription(codec);
	OH_AVFormat *avformat = OH_AVFormat_Create();
	OH_AVFormat_Copy(avformat, format);
	OH_AVFormat_Destroy(format);
	ohos::CodecData *codecData = encoder->GetCodecData();
	if (codecData != nullptr) {
		codecData->NeedInputBuffer(index, buffer);
	}
}

void OHOSMediaEncode::OnNewOutputBuffer(OH_AVCodec *codec, uint32_t index, OH_AVBuffer *buffer, void *userData) {
	(void)codec;
	if (userData == nullptr || buffer == nullptr) {
		return;
	}
	OHOSMediaEncode *encoder = static_cast<OHOSMediaEncode *>(userData);
	ohos::CodecData *codecData = encoder->GetCodecData();
	if (codecData != nullptr) {
		codecData->NewOutputBuffer(index, buffer);
	}
}

int32_t OHOSMediaEncode::PushInputData(
    uint8_t *pktData, size_t dataSize, int pts, int flag, std::chrono::milliseconds waitTimes) {
	OH_AVCodecBufferAttr attr = {0, 0, 0, AVCODEC_BUFFER_FLAGS_NONE};
	attr.pts = pts;
	attr.size = dataSize;
	attr.flags = flag;
	ohos::CodecData *codecData = data_.get();
	int indx = -1;
	int ret = codecData->InputData(pktData, dataSize, &attr, &indx, waitTimes);
	if (ret != ohos::ERRCODE_OK) {
		return ret;
	}
	return FreeInputData(indx);
}

int32_t OHOSMediaEncode::GetOutputBuffer(ohos::AvBufferInfo &info, std::chrono::milliseconds waitTimes) {
	return data_->GetOutputBuffer(info, waitTimes);
}

int32_t OHOSMediaEncode::GetInputBuffer(ohos::AvBufferInfo &info, std::chrono::milliseconds waitTimes) {
	return data_->GetInputBuffer(info, waitTimes);
}

int32_t OHOSMediaEncode::FreeInputData(uint32_t index) {
	int32_t ret = OH_VideoEncoder_PushInputBuffer(_impl, index);
	if (ret != AV_ERR_OK) {
		return AV_ERR_UNKNOWN;
	}
	return ret;
}

//填充
void OHOSMediaEncode::feed(mblk_t *rawData, uint64_t time, bool requestIFrame) {
	ms_message("OHOSMediaEncode: %s,%d", __FUNCTION__, __LINE__);
	std::unique_ptr<mblk_t, void (*)(mblk_t *)> rawDataPtr(rawData, freemsg);
	uint8_t *input_buf = NULL;
	if (_recoveryMode && time % 5000 == 0) {
		if (_impl == nullptr) createImpl();
		configureImpl();
		_recoveryMode = false;
	}

	if (_impl == nullptr) return;

	if (!_isRunning) {
		ms_error("MediaCodecEncoder: encoder not running. Dropping buffer.");
		return;
	}

	//  MSPicture pic;
	//  ms_yuv_buf_init_from_mblk(&pic, rawData);

	if (requestIFrame) {
		RequestIFrame(0);
		ms_error("MediaCodecEncoder: I-frame requested to MediaCodec");
	}
	size_t dataSize = rawData->b_wptr - rawData->b_rptr;

	ohos::AvBufferInfo inputBuffer{0, nullptr};
	OH_AVCodecBufferAttr attr;
	int ret = GetInputBuffer(inputBuffer, std::chrono::milliseconds(TIMEOUT));
	if (ret != ohos::ERRCODE_OK) {
		FreeInputData(inputBuffer.Index());
		return;
	}

	input_buf = inputBuffer.AvBufferAddr();
	if (input_buf == nullptr) {
		FreeInputData(inputBuffer.Index());
		return;
	}
	size_t bufferSize = inputBuffer.AvBufferSize();
	ms_error("OHOSMediaEncode: %s,%d , input 111buffer  bufferSize:%zu ;dataSize:%zu ", __FUNCTION__, __LINE__,
	         bufferSize, dataSize);
	if (bufferSize < dataSize) {
		dataSize = bufferSize;
	}
	memcpy(input_buf, rawData->b_rptr, dataSize);
	attr = *inputBuffer.Attr();
	attr.pts = time;
	attr.size = inputBuffer.AvBufferSize();
	attr.flags = AVCODEC_BUFFER_FLAGS_NONE;
	inputBuffer.Attr(&attr);
	FreeInputData(inputBuffer.Index());
	if (!_firstImageQueued) _firstImageQueued = true;
	_pendingFrames++;
	ms_message("OHOSMediaEncode: %s,%d end, dataSize---%zu", __FUNCTION__, __LINE__, dataSize);
}

//获取
bool OHOSMediaEncode::fetch(MSQueue *encodedData) {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	MSQueue outq;
	size_t bufsize;

	if (_impl == nullptr || !_isRunning || _recoveryMode || !_firstImageQueued) return false;

	if (_hasOutbufferDequeueLimit && _pendingFrames <= 0) return false;

	ms_queue_init(&outq);

	ohos::AvBufferInfo outputBuffer{0, nullptr};
	int ret = GetOutputBuffer(outputBuffer, std::chrono::milliseconds(TIMEOUT));
	if (ret != ohos::ERRCODE_OK) {
		return false;
	}

	OH_AVCodecBufferAttr attr = *outputBuffer.Attr();
	uint8_t *enc_frame_whole = outputBuffer.AvBufferAddr();
	bufsize = attr.size;
	_pendingFrames--;

	mediastreamer::H26xUtils::byteStreamToNalus(enc_frame_whole, bufsize, &outq, false);
	_psInserter->process(&outq, encodedData);
	ms_message("%s %s,%d---->bufsize:%zu", THIS_FILE, __FUNCTION__, __LINE__, bufsize);
	FreeOutputData(outputBuffer.Index());
	return true;
}

int32_t OHOSMediaEncode::FreeOutputData(uint32_t index) {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	int ret = OH_VideoEncoder_FreeOutputBuffer(_impl, index);
	if (ret != AV_ERR_OK) {
		ms_error("%s %s,%d errorCode", THIS_FILE, __FUNCTION__, __LINE__);
		return -1;
	}
	return 0;
}

void OHOSMediaEncode::configureImpl() {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	OH_AVFormat *format = createMediaFormat();
	int32_t ret = OH_VideoEncoder_Configure(_impl, format);
	if (ret != AV_ERR_OK) {
		ms_error("%s %s,%d errorCode", THIS_FILE, __FUNCTION__, __LINE__);
	}
	preConfig();
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	OH_AVFormat_Destroy(format);
}

void OHOSMediaEncode::RequestIFrame(int32_t value) {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	OH_AVFormat *avformat = OH_AVFormat_Create();
	if (avformat == nullptr) {
		ms_error("%s %s,%d errorCode", THIS_FILE, __FUNCTION__, __LINE__);
		return;
	}
	OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_REQUEST_I_FRAME, value);
	auto status = OH_VideoEncoder_SetParameter(_impl, avformat);
	if (status != AV_ERR_OK) {
		ms_error("%s %s,%d errorCode", THIS_FILE, __FUNCTION__, __LINE__);
		return;
	}
	OH_AVFormat_Destroy(avformat);
}

OH_AVFormat *OHOSMediaEncode::createMediaFormat() const {
	ms_message("%s %s,%d", THIS_FILE, __FUNCTION__, __LINE__);
	OH_AVFormat *avformat = OH_AVFormat_Create();
	OH_AVFormat_SetStringValue(avformat, OH_MD_KEY_CODEC_MIME, _mime.c_str());
	OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_PIXEL_FORMAT, AV_PIXEL_FORMAT_YUVI420);
	OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_HEIGHT, _vsize.height);
	OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_WIDTH, _vsize.width);
	OH_AVFormat_SetDoubleValue(avformat, OH_MD_KEY_FRAME_RATE, _fps);
	OH_AVFormat_SetLongValue(avformat, OH_MD_KEY_BITRATE, (_bitrate * 9) / 10);
	OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_VIDEO_ENCODE_BITRATE_MODE, _bitrateMode);
	OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_I_FRAME_INTERVAL, _iFrameInterval);
	OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_VIDEO_ENABLE_LOW_LATENCY, _encodingLatency);
	ms_message("%s %s,%d,OHOSCaptureVidoss createMediaFormat height:%d,width:%d", THIS_FILE, __FUNCTION__, __LINE__,
	           _vsize.height, _vsize.width);
	return avformat;
}

void OHOSMediaEncode::enableOutbufferDequeueLimit(bool enable) {
	if (enable) {
		ms_warning("Enabling the DEVICE_MCH265_LIMIT_DEQUEUE_OF_OUTPUT_BUFFERS hack on this device.");
	}
	_hasOutbufferDequeueLimit = enable;
}